#!/usr/bin/env python3
# coding: utf-8
# File: question_parser.py
# Author: lhy<lhy_in_blcu@126.com,https://huangyong.github.io>
# Date: 18-10-4
import ast
class QuestionPaser:
    '''构建实体节点'''

    def build_entitydict(self, args):
        entity_dict = {}
        for arg, types in args.items():
            for type in types:
                if type not in entity_dict:
                    entity_dict[type] = [arg]
                else:
                    entity_dict[type].append(arg)

        return entity_dict

    '''解析主函数'''

    def parser_main(self, res_classify):
        args = res_classify['args']
        entity_dict = self.build_entitydict(args)
        question_types = res_classify['question_types']
        sqls = []
        for question_type in question_types:
            sql_ = {}
            sql_['question_type'] = question_type
            sql = []
            # 简单问句
            if question_type == 'disease_symptom':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'symptom_disease':
                sql_['entities'] = entity_dict.get('symptom')
                sql = self.sql_transfer(question_type, entity_dict.get('symptom'))

            elif question_type == 'disease_cause':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_acompany':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_not_food':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_do_food':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'food_not_disease':
                sql_['entities'] = entity_dict.get('food')
                sql = self.sql_transfer(question_type, entity_dict.get('food'))

            elif question_type == 'food_do_disease':
                sql_['entities'] = entity_dict.get('food')
                sql = self.sql_transfer(question_type, entity_dict.get('food'))

            elif question_type == 'disease_drug':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'drug_disease':
                sql_['entities'] = entity_dict.get('drug')
                sql = self.sql_transfer(question_type, entity_dict.get('drug'))

            elif question_type == 'disease_check':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'check_disease':
                sql_['entities'] = entity_dict.get('check')
                sql = self.sql_transfer(question_type, entity_dict.get('check'))

            elif question_type == 'disease_prevent':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_lasttime':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_cureway':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_cureprob':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_easyget':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_desc':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'disease_department':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))

            elif question_type == 'department_disease':
                sql_['entities'] = entity_dict.get('department')
                sql = self.sql_transfer(question_type, entity_dict.get('department'))
            elif question_type == 'drug_producer':
                sql_['entities'] = entity_dict.get('drug')
                sql = self.sql_transfer(question_type, entity_dict.get('drug'))


            # 已知头实体和尾实体，求是否存在关系
            elif question_type == 'is_symptom':
                temp = entity_dict.get('disease').copy()
                temp.extend(entity_dict.get('symptom'))
                sql_['entities'] = list(set(temp))
                sql = self.sql_transfer(question_type, sql_['entities'])
            elif question_type == 'is_belonging':
                tuples = entity_dict.get('disease') + entity_dict.get('department')
                sql_['entities'] = tuples
                sql = self.sql_transfer(question_type, tuples)
            elif question_type == 'is_accompany_with':
                sql_['entities'] = entity_dict.get('disease')
                sql = self.sql_transfer(question_type, entity_dict.get('disease'))
            elif question_type == 'is_food':
                temp = entity_dict.get('disease').copy()
                temp.extend(entity_dict.get('food'))
                sql_['entities'] = temp
                sql = self.sql_transfer(question_type, temp)
            elif question_type == 'is_need_check':
                temp = entity_dict.get('disease').copy()
                temp.extend(entity_dict.get('check'))
                sql_['entities'] = temp
                sql = self.sql_transfer(question_type, temp)
            elif question_type == 'is_recommend_drug':
                temp = entity_dict.get('disease').copy()
                temp.extend(entity_dict.get('drug'))
                sql_['entities'] = temp
                sql = self.sql_transfer(question_type, temp)

            # 复杂问句
            # 两个疾病的相同症状、疾病治愈率比较、疾病治疗时间比较
            elif question_type in {'intersection_symptom', "compare_disease_cured_prob", "compare_disease_lasttime"}:
                temp = entity_dict.get('disease').copy()
                sql_['entities'] = temp
                sql = self.sql_transfer(question_type, temp)
            # 两个症状的相同疾病
            elif question_type == "intersection_disease":
                temp = entity_dict.get('symptom').copy()
                sql_['entities'] = temp
                sql = self.sql_transfer(question_type, temp)


            if sql:
                sql_['sql'] = sql

                sqls.append(sql_)

        return sqls

    '''针对不同的问题，分开进行处理'''

    def sql_transfer(self, question_type, entities):
        if not entities:
            return []

        # 查询语句
        sql = []

        # 查询疾病的原因
        if question_type == 'disease_cause':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.cause".format(i) for i in entities]

        # 查询疾病的防御措施
        elif question_type == 'disease_prevent':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.prevent".format(i) for i in entities]

        # 查询疾病的持续时间
        elif question_type == 'disease_lasttime':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.cure_lasttime".format(i) for i in entities]

        # 查询疾病的治愈概率
        elif question_type == 'disease_cureprob':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.cured_prob".format(i) for i in entities]

        # 查询疾病的治疗方式
        elif question_type == 'disease_cureway':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.cure_way".format(i) for i in entities]

        # 查询疾病的易发人群
        elif question_type == 'disease_easyget':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.easy_get".format(i) for i in entities]

        # 查询疾病的相关介绍
        elif question_type == 'disease_desc':
            sql = ["MATCH (m:Disease) where m.name = '{0}' return m.name, m.desc".format(i) for i in entities]

        # 查询疾病有哪些症状(包含KG）
        elif question_type == 'disease_symptom':
            sql = [
                "MATCH (m:Disease)-[r:has_symptom]->(n:Symptom) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]

        # 查询症状会导致哪些疾病
        elif question_type == 'symptom_disease':
            sql = [
                "MATCH (m:Disease)-[r:has_symptom]->(n:Symptom) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]

        # 查询疾病的并发症
        elif question_type == 'disease_acompany':
            sql = []
            for i in entities:
                sql1 = [
                "MATCH (m:Disease)-[r:acompany_with]->(n:Disease) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i)]
                sql2 = [
                "MATCH (m:Disease)-[r:acompany_with]->(n:Disease) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i)]
                sql += sql1 + sql2
        # 查询疾病的忌口
        elif question_type == 'disease_not_food':
            sql = ["MATCH (m:Disease)-[r:no_eat]->(n:Food) where m.name = '{0}' return m.name, r.name, n.name".format(i)
                   for i in entities]

        # 查询疾病建议吃的东西
        elif question_type == 'disease_do_food':
            sql1 = [
                "MATCH (m:Disease)-[r:do_eat]->(n:Food) where m.name = '{0}' return m.name, r.name, n.name".format(i)
                for i in entities]
            sql2 = [
                "MATCH (m:Disease)-[r:recommand_eat]->(n:Food) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]
            sql = sql1 + sql2

        # 已知忌口查疾病
        elif question_type == 'food_not_disease':
            sql = ["MATCH (m:Disease)-[r:no_eat]->(n:Food) where n.name = '{0}' return m.name, r.name, n.name".format(i)
                   for i in entities]

        # 已知推荐查疾病
        elif question_type == 'food_do_disease':
            sql1 = [
                "MATCH (m:Disease)-[r:do_eat]->(n:Food) where n.name = '{0}' return m.name, r.name, n.name".format(i)
                for i in entities]
            sql2 = [
                "MATCH (m:Disease)-[r:recommand_eat]->(n:Food) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]
            sql = sql1 + sql2

        # 查询疾病常用药品－药品别名记得扩充
        elif question_type == 'disease_drug':
            sql1 = [
                "MATCH (m:Disease)-[r:common_drug]->(n:Drug) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]
            sql2 = [
                "MATCH (m:Disease)-[r:recommand_drug]->(n:Drug) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]
            sql = sql1 + sql2

        # 已知药品查询能够治疗的疾病
        elif question_type == 'drug_disease':
            sql1 = [
                "MATCH (m:Disease)-[r:common_drug]->(n:Drug) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]
            sql2 = [
                "MATCH (m:Disease)-[r:recommand_drug]->(n:Drug) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]
            sql = sql1 + sql2

        # 已知药品种类查询在售药品
        elif question_type == 'drug_producer':
            sql = [
                "MATCH (m:Producer)-[r:drugs_of]->(n:Drug) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities
            ]

        # 查询疾病应该进行的检查
        elif question_type == 'disease_check':
            sql = [
                "MATCH (m:Disease)-[r:need_check]->(n:Check) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]

        # 已知检查查询疾病
        elif question_type == 'check_disease':
            sql = [
                "MATCH (m:Disease)-[r:need_check]->(n:Check) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities]

        # 已知疾病查相关科室
        elif question_type == 'disease_department':
            sql = [
                "MATCH (m:Disease)-[r:belongs_to]->(n:Department) where m.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities
            ]
        elif question_type == 'department_disease':
            sql = [
                "MATCH (m:Disease)-[r:belongs_to]->(n:Department) where n.name = '{0}' return m.name, r.name, n.name".format(
                    i) for i in entities
            ]
        #         '''
        #             已知头实体和尾实体，查询是否存在关系
        #         '''
        elif question_type == 'is_symptom':
            sql = [
                "MATCH (m)-[r:has_symptom]-(n) where m.name = '{0}' and n.name = '{1}' return labels(m), m.name , r.name,labels(n), n.name".format(
                    entities[0],
                    entities[1])]
        elif question_type == 'is_belonging':
            sql = [
                "MATCH (m:Disease)-[r:belongs_to]->(n:Department) WHERE m.name ='{0}' and n.name = '{1}' return m.name,r.name,n.name".format(
                    entities[0], entities[1])]
        elif question_type == 'is_accompany_with':
            sql = ["MATCH p = (m:Disease)-[r:acompany_with]-(n:Disease) WHERE m.name ='{0}' and n.name = '{1}' return p".format(
                    entities[0], entities[1])]
        elif question_type == 'is_food':
            sql = ["MATCH (m:Disease)-[r]->(n:Food) Where m.name ='{0}' and n.name = '{1}' return m.name,r.name,n.name".format(
                    entities[0], entities[1])]
        elif question_type == 'is_need_check':
            sql = ["MATCH (m:Disease)-[r:need_check]->(n:Check) Where m.name ='{0}' and n.name = '{1}' return m.name,r.name,n.name".format(
                    entities[0], entities[1])]
        elif question_type == 'is_recommend_drug':
            sql = ["MATCH (m:Disease)-[r]->(n:Drug) Where m.name ='{0}' and n.name = '{1}' return m.name,r.name,n.name".format(
                    entities[0], entities[1])]

        # '''
        #     复杂问句
        # '''
        # 询问两个疾病实体之间的相同症状
        elif question_type == 'intersection_symptom':
            sql = ["MATCH p=(m1)-[r1:has_symptom]->(n1),q=(m2)-[r2:has_symptom]->(n2) WHERE m1.name = '{0}' and m2.name = '{1}' and n1.name=n2.name RETURN n1.name, r1.name".format(
                    entities[0], entities[1])]
        # 询问同时具备两种症状的疾病
        elif question_type == 'intersection_disease':
            sql = ["MATCH p=(m1)-[r1:has_symptom]->(n1),q=(m2)-[r2:has_symptom]->(n2) WHERE n1.name = '{0}' and n2.name = '{1}' and m1.name=m2.name RETURN m1.name, r1.name".format(
                    entities[0], entities[1])]
        # 疾病治疗率比较
        elif question_type == 'compare_disease_cured_prob':
            sql = ["MATCH (m1:Disease), (m2:Disease) WHERE m1.name = '{0}' and m2.name = '{1}' RETURN m1.cured_prob, m2.cured_prob".format(
                    entities[0], entities[1])]
        # 疾病治疗周期比较
        elif question_type == 'compare_disease_lasttime':
            sql = ["MATCH (m1:Disease), (m2:Disease) WHERE m1.name = '{0}' and m2.name = '{1}' RETURN m1.cure_lasttime, m2.cure_lasttime".format(
                    entities[0], entities[1])]

        return sql


if __name__ == '__main__':
    handler = QuestionPaser()

    # res_sql = handler.parser_main(
    #     {'args': {'感冒': ['disease'], '脑膜炎': ['disease', 'symptom']}, 'question_types': ['is_symptom', 'is_accompany_with', 'disease_symptom', 'symptom_disease', 'disease_acompany']})
    # res_sql = handler.parser_main(
    #     {'args': {'毛囊炎': ['disease'], '银屑病': ['disease']},
    #      'question_types': ['compare_disease_lasttime']})
    while 1:
        question = input('input an question:')
        data = handler.parser_main(ast.literal_eval(question))
        print(data)
